home *** CD-ROM | disk | FTP | other *** search
/ MacHack 2000 / MacHack 2000.toast / pc / The Hacks / MacHacksBug / Python 1.5.2c1 / Extensions / Numerical / Lib / Numeric.py < prev    next >
Encoding:
Python Source  |  2000-06-23  |  9.5 KB  |  347 lines
  1. """Python code in support of array's, umath, and numeric 
  2. """
  3. __version__ = "11"
  4. __LLNLDistribution__ = __version__  
  5.  
  6. import multiarray
  7. from umath import * # Substitute fast_umath for "unsafe" math
  8. from Precision import *
  9.  
  10. import string, types, math
  11.  
  12. # make sure freeze knows _numpy is needed (it's only used by C code!)
  13. import _numpy
  14.  
  15.  
  16. #Use this to add a new axis to an array
  17. NewAxis = None
  18.  
  19. #The following functions are considered builtin, they all might be
  20. #in C some day
  21.  
  22. def arrayrange(start, stop=None, step=1, typecode=None):
  23.     """Just like range() except it returns a array whose
  24.     type can be specfied by the keyword argument typecode
  25.     """
  26.     
  27.     if (stop == None):
  28.         stop = start
  29.         start = 0
  30.     n = int(math.ceil(float(stop-start)/step))
  31.     if n <= 0:
  32.         m = zeros( (0,) )+(step+start+stop)
  33.     else:
  34.         m = (add.accumulate(ones((n,), Int))-1)*step +start+(stop-stop)
  35.         # the last bit is to deal with e.g. Longs -- 3L-3L==0L
  36.     if typecode != None and m.typecode() != typecode:
  37.         return m.astype(typecode)
  38.     else:
  39.         return m
  40.  
  41. #Include some functions straight from multiarray
  42. array = multiarray.array
  43. zeros = multiarray.zeros
  44. fromstring = multiarray.fromstring
  45. take = multiarray.take
  46. reshape = multiarray.reshape
  47. repeat = multiarray.repeat
  48. choose = multiarray.choose
  49. cross_correlate = multiarray.cross_correlate
  50. def convolve(a,v,mode=0):
  51.     if (len(v) > len(a)):
  52.         temp = a
  53.         a = v
  54.         v = temp
  55.         del temp
  56.     return cross_correlate(a,asarray(v)[::-1],mode)
  57.  
  58. ArrayType = multiarray.arraytype
  59.  
  60. def swapaxes(a, axis1, axis2):
  61.     n = len(shape(a))
  62.     if n <= 1: return a
  63.     new_axes = arange(n)
  64.     new_axes[axis1] = axis2
  65.     new_axes[axis2] = axis1
  66.     return multiarray.transpose(a, new_axes)
  67.  
  68. arraytype = multiarray.arraytype
  69. #add extra intelligence to the basic C functions
  70. def concatenate(a, axis=0):
  71.     if axis == 0:
  72.         return multiarray.concatenate(a)
  73.     else:
  74.         new_list = []
  75.         for m in a:
  76.             new_list.append(swapaxes(m, axis, 0))
  77.     return swapaxes(multiarray.concatenate(new_list), axis, 0)
  78.  
  79. def transpose(a, axes=None):
  80.     if axes == None:
  81.         axes = arange(len(array(a).shape))[::-1]
  82.     return multiarray.transpose(a, axes)
  83.  
  84. def sort(a, axis=-1):
  85.     if axis != -1: a = swapaxes(a, axis, -1)
  86.     s = multiarray.sort(a)
  87.     if axis != -1: s = swapaxes(s, axis, -1)
  88.     return s
  89.  
  90. def argsort(a, axis=-1):
  91.     if axis != -1: a = swapaxes(a, axis, -1)
  92.     s = multiarray.argsort(a)
  93.     if axis != -1: s = swapaxes(s, axis, -1)
  94.     return s
  95.  
  96. def argmax(a, axis=-1):
  97.     if axis != -1: a = swapaxes(a, axis, -1)
  98.     s = multiarray.argmax(a)
  99.     #probably need a swap here if > 2d
  100.     #if axis != -1: s = swapaxes(s, axis, -1)
  101.     return s
  102.  
  103. def argmin(x, axis=-1):
  104.     return argmax(negative(x), axis)
  105.  
  106.  
  107. searchsorted = multiarray.binarysearch
  108.  
  109. def innerproduct(a,b):
  110.     try:
  111.         return multiarray.innerproduct(a,b)
  112.     except(TypeError):
  113.         if array(a).shape == () or array(b).shape == ():
  114.             return a*b
  115.         else:
  116.             raise TypeError, "invalid types for dot"
  117.  
  118. def dot(a, b):
  119.     return innerproduct(a, swapaxes(b, -1, -2))
  120.  
  121. #This is obsolete, don't use in new code
  122. matrixmultiply = dot
  123.  
  124. #Use Konrad's printing function (modified for both str and repr now)
  125. from ArrayPrinter import array2string
  126. def array_repr(a, max_line_width = None, precision = None, suppress_small = None):
  127.     return array2string(a, max_line_width, precision, suppress_small, ', ', 1)
  128.  
  129. def array_str(a, max_line_width = None, precision = None, suppress_small = None):
  130.     return array2string(a, max_line_width, precision, suppress_small, ' ', 0)
  131.     
  132. multiarray.set_string_function(array_str, 0)
  133. multiarray.set_string_function(array_repr, 1)
  134.  
  135. #This is a nice value to have around
  136. #Maybe in sys some day
  137. LittleEndian = fromstring("\001"+"\000"*7, 'i')[0] == 1
  138.  
  139. def resize(a, new_shape):
  140.     """Returns a new array with the specified shape.  The original
  141.     array's total size can be any size.
  142.     """
  143.  
  144.     a = ravel(a)
  145.     if not len(a): return zeros(new_shape, a.typecode())
  146.     total_size = multiply.reduce(new_shape)
  147.     n_copies = total_size / len(a)
  148.     extra = total_size % len(a)
  149.  
  150.     if extra != 0: 
  151.         n_copies = n_copies+1
  152.         extra = len(a)-extra
  153.  
  154.     a = concatenate( (a,)*n_copies)
  155.     if extra > 0:
  156.         a = a[:-extra]
  157.  
  158.     return reshape(a, new_shape)
  159.  
  160. def indices(dimensions, typecode=None):
  161.     tmp = ones(dimensions, typecode)
  162.     lst = []
  163.     for i in range(len(dimensions)):
  164.         lst.append( add.accumulate(tmp, i, )-1 )
  165.     return array(lst)
  166.  
  167. def fromfunction(function, dimensions):
  168.     return apply(function, tuple(indices(dimensions)))
  169.     
  170.  
  171. def diagonal(a, offset= 0, axis1=0, axis2=1):
  172.     a = array (a)
  173.     if axis2 < axis1: axis1, axis2 = axis2, axis1
  174.     if axis2 > 1:
  175.         new_axes = range (len (a.shape))
  176.         del new_axes [axis2]; del new_axes [axis1]
  177.         new_axes [0:0] = [axis1, axis2]
  178.         a = transpose (a, new_axes)
  179.     s = a.shape
  180.     if len (s) == 2:
  181.         n1 = s [0]
  182.         n2 = s [1]
  183.         n = n1 * n2
  184.         s = (n,)
  185.         a = reshape (a, s)
  186.         if offset < 0:
  187.             return take (a, range ( - n2 * offset, min(n2, n1+offset) * (n2+1) - n2 * offset, n2+1), 0)
  188.         else:
  189.             return take (a, range (offset,         min(n1, n2-offset) * (n2+1) + offset,      n2+1), 0)
  190.     else :
  191.         my_diagonal = []
  192.         for i in range (s [0]) :
  193.             my_diagonal.append (diagonal (a [i], offset))
  194.         return array (my_diagonal)
  195.  
  196. def trace(a, offset=0, axis1=0, axis2=1):
  197.     return add.reduce(diagonal(a, offset, axis1, axis2))
  198.  
  199.  
  200. # These two functions are used in my modified pickle.py so that
  201. # matrices can be pickled.  Notice that matrices are written in 
  202. # binary format for efficiency, but that they pay attention to
  203. # byte-order issues for  portability.
  204.  
  205. def DumpArray(m, fp):    
  206.     if m.typecode() == 'O': 
  207.         raise TypeError, "Numeric Pickler can't pickle arrays of Objects"
  208.     s = m.shape
  209.     if LittleEndian: endian = "L"
  210.     else: endian = "B"
  211.     fp.write("A%s%s%d " % (m.typecode(), endian, m.itemsize()))
  212.     for d in s:
  213.         fp.write("%d "% d)
  214.     fp.write('\n')
  215.     fp.write(m.tostring())
  216.  
  217. def LoadArray(fp):
  218.     ln = string.split(fp.readline())
  219.     if ln[0][0] == 'A': ln[0] = ln[0][1:] # Nasty hack showing my ignorance of pickle
  220.     typecode = ln[0][0]
  221.     endian = ln[0][1]
  222.     
  223.     shape = map(lambda x: string.atoi(x), ln[1:])
  224.     itemsize = string.atoi(ln[0][2:])
  225.  
  226.     sz = reduce(multiply, shape)*itemsize
  227.     data = fp.read(sz)
  228.         
  229.     m = fromstring(data, typecode)
  230.     m = reshape(m, shape)
  231.  
  232.     if (LittleEndian and endian == 'B') or (not LittleEndian and endian == 'L'):
  233.         return m.byteswapped()
  234.     else:
  235.         return m
  236.  
  237. import pickle, copy
  238. class Unpickler(pickle.Unpickler):
  239.     def load_array(self):
  240.         self.stack.append(LoadArray(self))
  241.     
  242.     dispatch = copy.copy(pickle.Unpickler.dispatch)    
  243.     dispatch['A'] = load_array
  244.  
  245. class Pickler(pickle.Pickler):
  246.     def save_array(self, object):
  247.         DumpArray(object, self)
  248.  
  249.     dispatch = copy.copy(pickle.Pickler.dispatch)        
  250.     dispatch[ArrayType] = save_array
  251.  
  252. #Convenience functions
  253. from StringIO import StringIO
  254.  
  255. def dump(object, file):
  256.     Pickler(file).dump(object)
  257.  
  258. def dumps(object):
  259.     file = StringIO()
  260.     Pickler(file).dump(object)
  261.     return file.getvalue()
  262.  
  263. def load(file):
  264.     return Unpickler(file).load()
  265.  
  266. def loads(str):
  267.     file = StringIO(str)
  268.     return Unpickler(file).load()
  269.  
  270.  
  271. # These are all essentially abbreviations
  272. # These might wind up in a special abbreviations module
  273.  
  274. def ravel(m):
  275.     """Returns a 1d array corresponding to all the elements of it's
  276.     argument.
  277.     """
  278.     return reshape(m, (-1,))
  279.  
  280. def nonzero(a):
  281.     """Return the indices of the elements of a which are not zero, a must be 1d
  282.     """
  283.     return repeat(arange(len(a)), not_equal(a, 0))
  284.  
  285. def asarray(a, typecode=None):
  286.     return array(a, typecode, copy=0)
  287.  
  288. #Move this into C to do it right!
  289. def shape(a):
  290.     return asarray(a).shape
  291.  
  292. def where(condition, x, y):
  293.     """where(condition,x,y) is shaped like condition and has elements of x and
  294.     y where condition is respectively true or false
  295.     """
  296.     return choose(not_equal(condition, 0), (y, x))
  297.  
  298. def compress(condition, m, dimension=-1):
  299.     """compress(condition, x, dimension=-1) = those elements of x corresponding 
  300.     to those elements of condition that are "true".  condition must be the
  301.     same size as the given dimension of x."""
  302.     return take(m, nonzero(condition), dimension)
  303.  
  304. def clip(m, m_min, m_max):
  305.     """clip(m, m_min, m_max) = every entry in m that is less than m_min is
  306.     replaced by m_min, and every entry greater than m_max is replaced by
  307.     m_max.
  308.     """
  309.  
  310.     selector = less(m, m_min)+2*greater(m, m_max)
  311.     return choose(selector, (m, m_min, m_max))
  312.  
  313. def ones(shape, typecode='l'):
  314.     """ones(shape, typecode=Int) returns a array of the given dimensions
  315.     which is initialized to all ones.
  316.     """
  317.  
  318.     return zeros(shape, typecode)+array(1, typecode)
  319.  
  320. def identity(n):
  321.     return resize([1]+n*[0], (n,n))
  322.  
  323. sum = add.reduce
  324. cumsum = add.accumulate
  325. product = multiply.reduce
  326. cumproduct = multiply.accumulate
  327. alltrue = logical_and.reduce
  328. sometrue = logical_or.reduce
  329.  
  330. arange = arrayrange
  331.  
  332. ### Temporary solution for pickling arrays
  333. ### Quite inefficient
  334. ### david ascher, march 1998
  335. import copy_reg
  336.  
  337. def array_constructor(shape, typecode, thestr):
  338.     x = fromstring(thestr, typecode)
  339.     x.shape = shape
  340.     return x
  341.  
  342. def pickle_array(a):
  343.     return array_constructor, (a.shape, a.typecode(), a.tostring(),)  
  344.  
  345. copy_reg.pickle(ArrayType, pickle_array, array_constructor)
  346.  
  347.